# 生成随机参数结构，并运行FDTD计算
import os
import numpy as np
import lumapi as lum
import json as j
from time import time
from tqdm import tqdm
import threading


class ModPack:
    def __init__(self, filename, hide):
        self.p = None
        self.lx1 = None
        self.lx2 = None
        self.ly1 = None
        self.ly2 = None
        self.h1 = None
        self.h2 = None
        self.Rdata = None
        self.r = None
        self.gap = 60 * 1e-9
        self.filename = filename
        self.hide = hide
        self.mod = lum.FDTD(filename=filename, hide=self.hide)
        self.temp = ".\\temp_fdtd\\mod_copy_" + str(time()) + str(np.random.normal()) + ".fsp"
        self.temp_log = ".\\temp_fdtd\\mod_copy_" + str(time()) + str(np.random.normal()) + '_p0.log'
        # copy(filename, self.temp)
        self.mod.save(self.temp)

    def set_my_para(self, params):
        self.mod.switchtolayout()
        self.p = params['p']
        self.lx1 = params['lx1']
        self.lx2 = params['lx2']
        self.ly1 = params['ly1']
        self.ly2 = params['ly2']
        self.h1 = params['h1']
        self.h2 = params['h2']
        z_hmin = self.mod.getnamed("pillar::c1", "z min")
        self.mod.setnamed("pillar::c1", "y", 0.5 * (self.ly1 + self.gap))
        self.mod.setnamed("pillar::c2", "y", -0.5 * (self.ly2 + self.gap))
        self.mod.setnamed("pillar::c1", "x span", self.lx1)
        self.mod.setnamed("pillar::c2", "x span", self.lx2)
        self.mod.setnamed("pillar::c1", "y span", self.ly1)
        self.mod.setnamed("pillar::c2", "y span", self.ly2)
        self.mod.setnamed("FDTD", "x span", self.p)
        self.mod.setnamed("FDTD", "y span", self.p)
        self.mod.setnamed("pillar::c1", "z max", self.h1 + z_hmin)
        self.mod.setnamed("pillar::c2", "z max", self.h2 + z_hmin)
        self.mod.save()

    def run_my_mod(self):
        self.mod.switchtolayout()
        self.mod.run()
        self.r = self.mod.getresult("R", "T")
        self.Rdata = np.concatenate([self.r['lambda'], list(np.reshape(self.r['T'], (-1, 1)))], axis=1)

    def close_my_mod(self):
        self.mod.close()
        os.remove(self.temp)
        try:
            os.remove(self.temp_log)
        except:
            return None


# 产生随机结构参数，具有参数约束范围
def rand_para(gap=60):
    para_v = np.zeros(7)
    para_name = ['p', 'lx1', 'lx2', 'ly1', 'ly2', 'h1', 'h2']
    params = {}
    # p
    para_v[0] = np.random.randint(300, 400)
    # para_v[0] = 340
    # lx1
    para_v[1] = np.random.randint(60, para_v[0] - 60)
    # lx2
    para_v[2] = np.random.randint(60, para_v[0] - 60)
    # ly1
    para_v[3] = np.random.randint(50, para_v[0] - 2 * gap - 50)
    # ly2
    para_v[4] = np.random.randint(50, para_v[0] - (2 * gap) - para_v[3])
    # h1
    para_v[5] = np.random.randint(100, 500)
    # h2
    para_v[6] = np.random.randint(100, 500)
    for j in range(len(para_v)):
        params[para_name[j]] = para_v[j] * 1e-9

    return params


# 进行模拟，其中生成随机结构参数，返回参数和反射谱组成的元组列表
def sim_struc(name, n_batch, path=None, hide=False):
    if path is None:
        path = './TiO2-gen.fsp'
    my_sim = ModPack(path, hide)
    out_all = []
    struc_para = []
    R_data = []
    # print('Process: %s 开始计算' % name)
    with tqdm(total=int(n_batch), desc=name) as t:
        for num in range(int(n_batch)):
            init_t = time()
            pa = rand_para()
            my_sim.set_my_para(params=pa)
            my_sim.run_my_mod()

            struc_para.append(list(pa.values()))
            R_data.append(my_sim.Rdata.tolist())
            out = {'para_array': struc_para, 'R_data': R_data}
            with open('./result_mp/{0}.json'.format(name), 'w') as f:
                j.dump(out, f)
            # finished = 100 * ((num + 1) / n_batch)
            # dt = time() - init_t
            # print('%s第%d次计算完成，进度%.2f%%, 用时: %.2fs' % (name, num + 1, finished, dt))
            t.update()
    # print('\n---------Process %s finished---------' % name)

    out = {'para_array': struc_para, 'R_data': list(R_data)}
    my_sim.close_my_mod()
    return out


# 模拟任意结构参数
def sim_result(para_arr, hide):
    mod = ModPack('TiO2-gen.fsp', hide=hide)
    para_name = ['p', 'lx1', 'lx2', 'ly1', 'ly2', 'h1', 'h2']
    R_da = []
    for j in range(para_arr.shape[0]):
        para_j = dict(zip(para_name, list(para_arr[j, :])))
        mod.set_my_para(para_j)
        mod.run_my_mod()
        R_da_j = mod.Rdata
        R_da.append(R_da_j.tolist())
    mod.close_my_mod()
    return R_da
